from fastapi import FastAPI, HTTPException, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse  # 添加HTMLResponse导入
from pydantic import BaseModel
import requests
from PIL import Image
import torch
from torchvision import transforms, models
import numpy as np
import logging
import io
import os

app = FastAPI()
logger = logging.getLogger(__name__)

# 添加CORS中间件以支持跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应更具体
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化模型
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
model.eval()

async def process_image(image_url: str = None, upload_file: UploadFile = None):
    try:
        if image_url:
            # 禁用代理
            session = requests.Session()
            session.trust_env = False  # 不信任环境变量中的代理设置
            
            response = session.get(image_url, timeout=10)
            response.raise_for_status()
            image = Image.open(io.BytesIO(response.content)).convert('RGB')
            origin_type = 'url'
        elif upload_file:
            contents = await upload_file.read()
            image = Image.open(io.BytesIO(contents)).convert("RGB")
            origin_type = 'upload'

        input_tensor = preprocess(image)
        input_batch = input_tensor.unsqueeze(0)
        
        with torch.no_grad():
            output = model(input_batch)
        
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        top5_prob, top5_catid = torch.topk(probabilities, 5)
        
        return {
            "origin_type": origin_type,
            "predictions": [
                {
                    "label": models.ResNet50_Weights.IMAGENET1K_V1.meta["categories"][catid],
                    "confidence": f"{prob.item():.4f}" if isinstance(prob, torch.Tensor) else f"{prob:.4f}"
                } for catid, prob in zip(top5_catid, top5_prob)
            ]
        }
    except Exception as e:
        logger.error(f"处理错误: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


# 图像预处理流程
def preprocess(image):
    return transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225])
    ])(image)

class ImageRequest(BaseModel):
    image_url: str

@app.post("/predict")
async def predict_url(request: ImageRequest):
    """URL识别接口"""
    return await process_image(image_url=request.image_url)

@app.post("/upload")
async def upload_image(file: UploadFile = File(...)):
    return await process_image(upload_file=file)

# 添加静态文件服务
@app.get("/", response_class=HTMLResponse)
async def read_index():
    try:
        with open("index.html", "r", encoding="utf-8") as f:
            return HTMLResponse(f.read())
    except FileNotFoundError:
        return HTMLResponse(
            content="<h1>前端页面未找到</h1><p>请确保index.html文件存在于当前目录</p>",
            status_code=404
        )

# 健康检查端点
@app.get("/health")
async def health_check():
    return JSONResponse(content={"status": "healthy", "message": "服务正常运行"})

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)